###############################################################################
# BSD 3-Clause License
#
# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
#   
# Copyright (c) 2017, Soumith Chintala. All rights reserved.
###############################################################################
'''
Code adapted from https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py
Introduced partial convolutoins based padding for convolutional layers
'''

import torch.nn as nn
import math
import torch.utils.model_zoo as model_zoo
from .partialconv2d import PartialConv2d

__all__ = ['PDResNet', 'pdresnet18', 'pdresnet34', 'pdresnet50', 'pdresnet101',
           'pdresnet152']


# model_urls = {
#     'pdresnet18': '',
#     'pdresnet34': '',
#     'pdresnet50': '',
#     'pdresnet101': '',
#     'pdresnet152': '',
# }



def conv3x3(in_planes, out_planes, stride=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False)


class BasicBlock(nn.Module):
    expansion = 1

    def __init__(self, inplanes, planes, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = conv3x3(inplanes, planes, stride)
      
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
      
      
        self.stride = stride

    def forward(self, x):
        residual = x
        out= self.conv1(x)
        out = self.relu(out)
        out= self.conv2(out)
        out += residual
        out = self.relu(out)

        return out

def myresnet(input_channel=3):
    """Constructs a PDResNet-18 model.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
    """
    model=nn.Sequential(
        nn.Conv2d(input_channel, 64, kernel_size=3, padding=1),
        nn.ReLU(inplace=True),
        BasicBlock(64,64),
        BasicBlock(64,64),
        nn.Conv2d(64, 3, kernel_size=3, padding=1),
    )
    return model